{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Matrix Diagonalization\n",
    "Consider a $n \\times n$ matrix $A_{n, n}$ with $n$ linearly independent eigenvectors.\n",
    "<br> Let $S_{n,n}$ be a matrix with these eigenvectors as its columns.\n",
    "<br>Then $A_{n,n}$ can be factorized as\n",
    "$$\n",
    "A = S \\Sigma S^{-1}\n",
    "$$\n",
    "where $\\Sigma$ is a diagonal matrix (all off-digonal elements are zero).\n",
    "<br> The above matrix factorization is known as diagonalization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A =\n",
      "tensor([[ 0.7071,  0.7071,  0.0000],\n",
      "        [-0.7071,  0.7071,  0.0000],\n",
      "        [ 0.0000,  0.0000,  1.0000]])\n",
      "\n",
      "S =\n",
      "tensor([[0.7071+0.0000j, 0.7071-0.0000j, 0.0000+0.0000j],\n",
      "        [0.0000+0.7071j, 0.0000-0.7071j, 0.0000+0.0000j],\n",
      "        [0.0000+0.0000j, 0.0000-0.0000j, 1.0000+0.0000j]])\n",
      "\n",
      "sigma =\n",
      "tensor([[0.7071+0.7071j, 0.0000+0.0000j, 0.0000+0.0000j],\n",
      "        [0.0000+0.0000j, 0.7071-0.7071j, 0.0000+0.0000j],\n",
      "        [0.0000+0.0000j, 0.0000+0.0000j, 1.0000+0.0000j]])\n",
      "\n",
      "S_inv =\n",
      "tensor([[0.7071+0.0000j, 0.0000-0.7071j, 0.0000+0.0000j],\n",
      "        [0.7071+0.0000j, 0.0000+0.7071j, 0.0000+0.0000j],\n",
      "        [0.0000+0.0000j, 0.0000+0.0000j, 1.0000+0.0000j]])\n",
      "\n",
      "S sigmaa S_inv =\n",
      "tensor([[ 0.7071-5.5519e-08j,  0.7071+5.5519e-08j,  0.0000+0.0000e+00j],\n",
      "        [-0.7071-5.5519e-08j,  0.7071-5.5519e-08j,  0.0000+0.0000e+00j],\n",
      "        [ 0.0000+0.0000e+00j,  0.0000+0.0000e+00j,  1.0000+0.0000e+00j]])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "# Let us reconsider our rotation matrix\n",
    "def diagonalise(matrix):\n",
    "    try:\n",
    "        # np.linalg.eig(matrix) returns the eigenvalues of matrix in an\n",
    "        # array (first return value) and the eigvectors as a\n",
    "        # matrix (each column is an eigenvector)\n",
    "        l, e = torch.linalg.eig(matrix)\n",
    "        \n",
    "        # make a diagonal matrix from the eigenvalues\n",
    "        sigma = torch.diag(l)\n",
    "        \n",
    "        # return the three factor matrices\n",
    "        return e, torch.diag(l), torch.linalg.inv(e)\n",
    "    except RuntimeError:\n",
    "        print(\"Cannot diagonalise matrix!\")\n",
    "\n",
    "    \n",
    "A = torch.tensor([[0.7071, 0.7071, 0], [-0.7071, 0.7071, 0],\n",
    "              [0, 0, 1]])\n",
    "print(\"A =\\n{}\".format(A))\n",
    "S, sigma, S_inv = diagonalise(A)\n",
    "\n",
    "# Let us reconstruct the original matriox from its factors\n",
    "A1 = torch.matmul(S, torch.matmul(sigma, S_inv))\n",
    "print(\"\\nS =\\n{}\".format(S))\n",
    "print(\"\\nsigma =\\n{}\".format(sigma))\n",
    "print(\"\\nS_inv =\\n{}\".format(S_inv))\n",
    "print(\"\\nS sigmaa S_inv =\\n{}\".format(A1))\n",
    "\n",
    "# We assert that the original matrix is the same as the\n",
    "# reconstruction from the diagonal decomposition factors.\n",
    "\n",
    "assert torch.allclose(A, A1.real)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Solving Linear Systems without Inverse\n",
    "\n",
    "Let us try solving the following set of equations:\n",
    "\n",
    "\\begin{align*}\n",
    "    x + 2y + z    &= 8 \\\\\n",
    "    2x + 2y + 3z &= 15 \\\\\n",
    "    x + 3y + 3z  &= 16\n",
    "\\end{align*}\n",
    "\n",
    "This can be written using matrices and vectors as\n",
    "\\begin{equation*}\n",
    "A\\vec{x} = \\vec{b}\n",
    "\\end{equation*}\n",
    "where $A=\n",
    "\\begin{bmatrix}\n",
    "        1 & 2 & 1 \\\\\n",
    "        2 & 2 & 3 \\\\\n",
    "        1 & 3 & 3 \n",
    "\\end{bmatrix}\n",
    "\\;\\;\\;\\;\\;\\;\n",
    "\\vec{x} = \n",
    "\\begin{bmatrix}\n",
    "        x \\\\\n",
    "        y \\\\\n",
    "        z \n",
    "\\end{bmatrix}\n",
    "\\;\\;\\;\\;\\;\\;\n",
    "\\vec{b} = \n",
    "\\begin{bmatrix}\n",
    "        8 \\\\\n",
    "        15 \\\\\n",
    "        16\n",
    "\\end{bmatrix}$\n",
    "\n",
    "\n",
    "Note that $A$ is a symmetric matrix. It has orthogonal eigenvectros.\n",
    "<br> The matrix with eigenvectors of $A$ in columns is orthogonal.\n",
    "<br> Its transpose and inverse are same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Solution using inverse: tensor([1.0000, 2.0000, 3.0000])\n",
      "Solution using diagonalisation: tensor([1.0000+0.j, 2.0000+0.j, 3.0000+0.j])\n"
     ]
    }
   ],
   "source": [
    "A = torch.tensor([[1, 2, 1], [2, 2, 3], [1, 3, 3]], dtype=torch.float)\n",
    "assert torch.all(A == A.T)\n",
    "b = torch.tensor([8, 15, 16], dtype=torch.cfloat)\n",
    "\n",
    "# One way to solve for this is to compute matrix inverse\n",
    "# i.e x = A_inv b\n",
    "x_0 = torch.matmul(torch.linalg.inv(A), b.real)\n",
    "print(\"Solution using inverse: {}\".format(x_0))\n",
    "\n",
    "# Matrix inversion is a complex process that can be\n",
    "# numerically unstable. If possible we use diagonalisation.\n",
    "w, S = torch.linalg.eig(A)\n",
    "\n",
    "# We know that A = S sigma S_inv\n",
    "# So S sigma S_inv x = b\n",
    "\n",
    "# sigma S_inv x =  S_inv b  ==> sigma S_inv x = S_t b\n",
    "S_inv_b = torch.matmul(S.T, b)\n",
    "\n",
    "# => S_inv x = sigma_inv(S_t b)\n",
    "# Since inversion of the diagonal matrix is just division\n",
    "# of all elements by 1, we can compute \n",
    "# sigma_inv as follows\n",
    "sigma_inv = torch.diag(1/ w)\n",
    "sigma_inv_S_inv_b = torch.matmul(sigma_inv, S_inv_b)\n",
    "\n",
    "# => x = S (sigma_inv(S_t b))\n",
    "x_1 = torch.matmul(S, sigma_inv_S_inv_b) \n",
    "print(\"Solution using diagonalisation: {}\".format(x_1))\n",
    "\n",
    "assert torch.allclose(x_0, x_1.real)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Matrix powers using diagonalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A:\n",
      "tensor([[1., 2., 1.],\n",
      "        [2., 2., 3.],\n",
      "        [1., 3., 3.]])\n",
      "A**2 computed via brute force:\n",
      "tensor([[ 6.,  9., 10.],\n",
      "        [ 9., 17., 17.],\n",
      "        [10., 17., 19.]])\n",
      "A**2 computed using diagonalisation:\n",
      "tensor([[ 6.0000+0.j,  9.0000+0.j, 10.0000+0.j],\n",
      "        [ 9.0000+0.j, 17.0000+0.j, 17.0000+0.j],\n",
      "        [10.0000+0.j, 17.0000+0.j, 19.0000+0.j]])\n"
     ]
    }
   ],
   "source": [
    "A = torch.tensor([[1, 2, 1], [2, 2, 3], [1, 3, 3]], dtype=torch.float)\n",
    "\n",
    "def brute_force_matrix_power(matrix, y):\n",
    "    \"\"\"Returns matrix raised to the power of y\"\"\"\n",
    "    # Assert that it is a square matrix\n",
    "    assert len(matrix.shape) == 2\\\n",
    "           and matrix.shape[0] == matrix.shape[1] \n",
    "    assert type(y) == int \n",
    "    if y == 0:\n",
    "        # Return identity matrix \n",
    "        return torch.eye(matrix.shape[0])\n",
    "    output_matrix = torch.clone(matrix)\n",
    "    for i in range(y-1):\n",
    "        output_matrix = torch.matmul(matrix, output_matrix)\n",
    "    return output_matrix\n",
    "\n",
    "def matrix_power_using_diag(matrix, y):\n",
    "    \"\"\"Returns matrix raised to the power of y\"\"\"\n",
    "    # Assert that it is a square matrix\n",
    "    assert len(matrix.shape) == 2\\\n",
    "           and matrix.shape[0] == matrix.shape[1] \n",
    "    assert type(y) == int \n",
    "    if y == 0:\n",
    "        # Return identity matrix \n",
    "        return torch.eye(matrix.shape[0])\n",
    "    w, v = torch.linalg.eig(matrix)\n",
    "    w_y = w ** float(y) # Compute w^y\n",
    "    return torch.matmul(v, torch.matmul(torch.diag(w_y),\n",
    "                                  torch.linalg.inv(v)))\n",
    "    \n",
    "# One way to compute A**2 would be to multiply A by A\n",
    "# This is an O(nm^3) operation\n",
    "A_2_brute_force = brute_force_matrix_power(A, 2)\n",
    "\n",
    "\n",
    "# We can instead diagonalise the matrix and compute A ** 2\n",
    "# This reduces the complexity to O(m)\n",
    "A_2_diag = matrix_power_using_diag(A, 2)\n",
    "\n",
    "print(\"A:\\n{}\".format(A))\n",
    "print(\"A**2 computed via brute force:\\n{}\".\\\n",
    "      format(A_2_brute_force))\n",
    "print(\"A**2 computed using diagonalisation:\\n{}\".\\\n",
    "      format(A_2_diag))\n",
    "\n",
    "assert torch.allclose(A_2_brute_force, A_2_diag.real)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
